import argparse
from osgeo.gdal import Open, GetDriverByName, GDT_Int32


class Debris:
    def __init__(self, args, **kwargs):
        self.dem_path = None
        self.filtered_path = None
        self.output_path = None
        if args is not None:
            for key in args.__dict__:
                self[key] = args.__dict__[key]
        for key in kwargs.keys():
            self[key] = kwargs[key]

    def __getitem__(self, item):
        return self.__dict__[item]

    def __setitem__(self, key, value):
        self.__dict__[key] = value

    def __str__(self):
        return f"dem={self.dem_path},filtered={self.filtered_path},output={self.output_path}"

    def run(self):
        """
        执行杂物计算。

        :return:
        """
        # 加载输入数据
        dem = Open(self.dem_path)
        proj = dem.GetProjection()
        trans = dem.GetGeoTransform()
        n_rows, n_cols = dem.RasterYSize, dem.RasterXSize
        nodata = dem.GetRasterBand(1).GetNoDataValue()
        filtered = Open(self.filtered_path)

        dem_arr = dem.ReadAsArray()
        filtered_arr = filtered.ReadAsArray()
        res = dem_arr - filtered_arr
        # 若dem或filter的对应位置为nodata，结果均设置为0
        res[dem_arr == nodata] = 0
        res[filtered_arr == nodata] = 0
        # 以dem为界，将范围外的区域设置为nodata
        res[dem_arr == nodata] = -9999

        # 保存输出结果
        driver = GetDriverByName("GTiff")
        img_out = driver.Create(self.output_path, n_cols, n_rows, 1, GDT_Int32)
        img_out.SetGeoTransform(trans)
        img_out.SetProjection(proj)
        img_out.GetRasterBand(1).WriteArray(res)
        img_out.GetRasterBand(1).SetNoDataValue(-9999)
        del img_out


if __name__ == '__main__':
    # 接收控制台传入参数
    parser = argparse.ArgumentParser(description='classification parameters')
    parser.add_argument('dem_path', type=str, help='dem file path')
    parser.add_argument('filtered_path', type=str, help='filtered dem file path')
    parser.add_argument('output_path', type=str, help='output file path')
    args = parser.parse_args()
    # print(args.__dict__)
    dbs = Debris(args=args)
    dbs.run()
    # print(slp)
